"""
    Interpolation

Comprehensive interpolation methods for observation operators in GSI.

This module provides:
- Bilinear interpolation for 2D horizontal fields
- Trilinear interpolation for 3D spatial fields
- Vertical interpolation with various coordinate systems
- Adjoint operators for all interpolation methods
- Edge case handling and extrapolation strategies

# Mathematical Framework

For an observation at location (lon, lat, lev), interpolation computes:
    y = Σᵢ wᵢ · xᵢ

where wᵢ are interpolation weights and xᵢ are grid point values.

The adjoint operation distributes observation residuals back to grid points:
    ∂J/∂xᵢ = wᵢ · (y_obs - y_model)

# Adjoint Consistency

All interpolation methods satisfy the adjoint test:
    ⟨H(x), y⟩ = ⟨x, Hᵀ(y)⟩

where H is the forward interpolation operator and Hᵀ is its adjoint.
"""
module Interpolation

using LinearAlgebra
using StaticArrays

export InterpolationWeights, InterpolationMethod
export bilinear_interpolate, bilinear_interpolate_adjoint
export trilinear_interpolate, trilinear_interpolate_adjoint
export vertical_interpolate, vertical_interpolate_adjoint
export find_grid_cell, compute_bilinear_weights, compute_trilinear_weights
export pressure_to_sigma, sigma_to_pressure
export logarithmic_vertical_interpolate

# Data structures

"""
    InterpolationMethod

Enumeration of available interpolation methods.
"""
@enum InterpolationMethod begin
    NEAREST_NEIGHBOR
    BILINEAR
    TRILINEAR
    CUBIC
    SPLINE
end

"""
    InterpolationWeights{T<:AbstractFloat}

Stores precomputed interpolation weights and indices for efficient reuse.

# Fields
- `indices::Vector{CartesianIndex{N}}`: Grid point indices
- `weights::Vector{T}`: Corresponding interpolation weights
- `in_bounds::Bool`: Whether observation is within grid domain
"""
struct InterpolationWeights{T<:AbstractFloat,N}
    indices::Vector{CartesianIndex{N}}
    weights::Vector{T}
    in_bounds::Bool

    function InterpolationWeights{T,N}(indices::Vector{CartesianIndex{N}},
                                       weights::Vector{T},
                                       in_bounds::Bool=true) where {T,N}
        @assert length(indices) == length(weights) "Indices and weights must have same length"
        new{T,N}(indices, weights, in_bounds)
    end
end

# 2D Bilinear Interpolation

"""
    find_grid_cell(x::AbstractVector{T}, x_obs::T) -> Tuple{Int,Int,T}

Find grid cell containing observation point.

# Arguments
- `x`: Grid coordinate array (monotonically increasing or decreasing)
- `x_obs`: Observation coordinate

# Returns
- `(i1, i2, frac)`: Lower index, upper index, fractional position (0 to 1)
"""
function find_grid_cell(x::AbstractVector{T}, x_obs::T) where T<:AbstractFloat
    n = length(x)

    # Handle edge cases
    if x_obs <= x[1]
        return (1, 1, T(0))
    elseif x_obs >= x[end]
        return (n, n, T(1))
    end

    # Binary search for monotonically increasing array
    if x[end] > x[1]
        i1 = searchsortedlast(x, x_obs)
        i2 = min(i1 + 1, n)
    else  # Decreasing array
        i1 = searchsortedfirst(x, x_obs, rev=true)
        i2 = min(i1 + 1, n)
    end

    # Compute fractional position
    if i1 < n && abs(x[i2] - x[i1]) > eps(T)
        frac = (x_obs - x[i1]) / (x[i2] - x[i1])
    else
        frac = T(0)
    end

    return (i1, i2, clamp(frac, T(0), T(1)))
end

"""
    compute_bilinear_weights(lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                            lon_obs::T, lat_obs::T) -> InterpolationWeights{T,2}

Compute bilinear interpolation weights for 2D horizontal interpolation.

# Arguments
- `lon_grid`: Grid longitude coordinates
- `lat_grid`: Grid latitude coordinates
- `lon_obs`: Observation longitude
- `lat_obs`: Observation latitude

# Returns
- `InterpolationWeights` containing 4 grid points and weights
"""
function compute_bilinear_weights(lon_grid::AbstractVector{T},
                                 lat_grid::AbstractVector{T},
                                 lon_obs::T, lat_obs::T) where T<:AbstractFloat
    # Find grid cell in each dimension
    i1, i2, wx = find_grid_cell(lon_grid, lon_obs)
    j1, j2, wy = find_grid_cell(lat_grid, lat_obs)

    # Check if in bounds
    in_bounds = (i1 > 0 && i2 <= length(lon_grid) &&
                 j1 > 0 && j2 <= length(lat_grid))

    # Compute bilinear weights (4 points)
    w11 = (T(1) - wx) * (T(1) - wy)  # Lower-left
    w21 = wx * (T(1) - wy)            # Lower-right
    w12 = (T(1) - wx) * wy            # Upper-left
    w22 = wx * wy                      # Upper-right

    # Create indices and weights arrays
    indices = [CartesianIndex(i1, j1),
               CartesianIndex(i2, j1),
               CartesianIndex(i1, j2),
               CartesianIndex(i2, j2)]

    weights = [w11, w21, w12, w22]

    return InterpolationWeights{T,2}(indices, weights, in_bounds)
end

"""
    bilinear_interpolate(field::AbstractMatrix{T},
                        lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                        lon_obs::T, lat_obs::T) -> T

Perform bilinear interpolation on a 2D field.

# Arguments
- `field`: 2D field to interpolate (nlat × nlon)
- `lon_grid`: Longitude coordinates
- `lat_grid`: Latitude coordinates
- `lon_obs`: Observation longitude
- `lat_obs`: Observation latitude

# Returns
- Interpolated value at observation location
"""
function bilinear_interpolate(field::AbstractMatrix{T},
                             lon_grid::AbstractVector{T},
                             lat_grid::AbstractVector{T},
                             lon_obs::T, lat_obs::T) where T<:AbstractFloat
    weights = compute_bilinear_weights(lon_grid, lat_grid, lon_obs, lat_obs)

    if !weights.in_bounds
        return T(NaN)
    end

    # Compute weighted sum
    value = zero(T)
    for (idx, w) in zip(weights.indices, weights.weights)
        if checkbounds(Bool, field, idx)
            value += w * field[idx]
        end
    end

    return value
end

"""
    bilinear_interpolate_adjoint!(gradient::AbstractMatrix{T}, residual::T,
                                  lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                                  lon_obs::T, lat_obs::T)

Adjoint of bilinear interpolation - distributes residual back to grid points.

# Arguments
- `gradient`: Gradient array to accumulate into (modified in-place)
- `residual`: Observation residual to distribute
- `lon_grid`: Longitude coordinates
- `lat_grid`: Latitude coordinates
- `lon_obs`: Observation longitude
- `lat_obs`: Observation latitude
"""
function bilinear_interpolate_adjoint!(gradient::AbstractMatrix{T},
                                      residual::T,
                                      lon_grid::AbstractVector{T},
                                      lat_grid::AbstractVector{T},
                                      lon_obs::T, lat_obs::T) where T<:AbstractFloat
    weights = compute_bilinear_weights(lon_grid, lat_grid, lon_obs, lat_obs)

    if !weights.in_bounds
        return nothing
    end

    # Distribute residual according to interpolation weights (adjoint operation)
    for (idx, w) in zip(weights.indices, weights.weights)
        if checkbounds(Bool, gradient, idx)
            gradient[idx] += w * residual
        end
    end

    return nothing
end

# 3D Trilinear Interpolation

"""
    compute_trilinear_weights(lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                             lev_grid::AbstractVector{T},
                             lon_obs::T, lat_obs::T, lev_obs::T) -> InterpolationWeights{T,3}

Compute trilinear interpolation weights for 3D spatial interpolation.

# Arguments
- `lon_grid`: Grid longitude coordinates
- `lat_grid`: Grid latitude coordinates
- `lev_grid`: Grid vertical level coordinates
- `lon_obs`: Observation longitude
- `lat_obs`: Observation latitude
- `lev_obs`: Observation vertical level

# Returns
- `InterpolationWeights` containing 8 grid points and weights
"""
function compute_trilinear_weights(lon_grid::AbstractVector{T},
                                  lat_grid::AbstractVector{T},
                                  lev_grid::AbstractVector{T},
                                  lon_obs::T, lat_obs::T, lev_obs::T) where T<:AbstractFloat
    # Find grid cell in each dimension
    i1, i2, wx = find_grid_cell(lon_grid, lon_obs)
    j1, j2, wy = find_grid_cell(lat_grid, lat_obs)
    k1, k2, wz = find_grid_cell(lev_grid, lev_obs)

    # Check if in bounds
    in_bounds = (i1 > 0 && i2 <= length(lon_grid) &&
                 j1 > 0 && j2 <= length(lat_grid) &&
                 k1 > 0 && k2 <= length(lev_grid))

    # Compute trilinear weights (8 corner points of cube)
    wx1 = T(1) - wx
    wx2 = wx
    wy1 = T(1) - wy
    wy2 = wy
    wz1 = T(1) - wz
    wz2 = wz

    # 8 weights for 8 corners
    weights = [
        wx1 * wy1 * wz1,  # (i1, j1, k1)
        wx2 * wy1 * wz1,  # (i2, j1, k1)
        wx1 * wy2 * wz1,  # (i1, j2, k1)
        wx2 * wy2 * wz1,  # (i2, j2, k1)
        wx1 * wy1 * wz2,  # (i1, j1, k2)
        wx2 * wy1 * wz2,  # (i2, j1, k2)
        wx1 * wy2 * wz2,  # (i1, j2, k2)
        wx2 * wy2 * wz2   # (i2, j2, k2)
    ]

    indices = [
        CartesianIndex(i1, j1, k1),
        CartesianIndex(i2, j1, k1),
        CartesianIndex(i1, j2, k1),
        CartesianIndex(i2, j2, k1),
        CartesianIndex(i1, j1, k2),
        CartesianIndex(i2, j1, k2),
        CartesianIndex(i1, j2, k2),
        CartesianIndex(i2, j2, k2)
    ]

    return InterpolationWeights{T,3}(indices, weights, in_bounds)
end

"""
    trilinear_interpolate(field::AbstractArray{T,3},
                         lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                         lev_grid::AbstractVector{T},
                         lon_obs::T, lat_obs::T, lev_obs::T) -> T

Perform trilinear interpolation on a 3D field.

# Arguments
- `field`: 3D field to interpolate (nlon × nlat × nlev)
- `lon_grid`: Longitude coordinates
- `lat_grid`: Latitude coordinates
- `lev_grid`: Vertical level coordinates
- `lon_obs`: Observation longitude
- `lat_obs`: Observation latitude
- `lev_obs`: Observation vertical level

# Returns
- Interpolated value at observation location
"""
function trilinear_interpolate(field::AbstractArray{T,3},
                              lon_grid::AbstractVector{T},
                              lat_grid::AbstractVector{T},
                              lev_grid::AbstractVector{T},
                              lon_obs::T, lat_obs::T, lev_obs::T) where T<:AbstractFloat
    weights = compute_trilinear_weights(lon_grid, lat_grid, lev_grid,
                                       lon_obs, lat_obs, lev_obs)

    if !weights.in_bounds
        return T(NaN)
    end

    # Compute weighted sum
    value = zero(T)
    for (idx, w) in zip(weights.indices, weights.weights)
        if checkbounds(Bool, field, idx)
            value += w * field[idx]
        end
    end

    return value
end

"""
    trilinear_interpolate_adjoint!(gradient::AbstractArray{T,3}, residual::T,
                                   lon_grid::AbstractVector{T}, lat_grid::AbstractVector{T},
                                   lev_grid::AbstractVector{T},
                                   lon_obs::T, lat_obs::T, lev_obs::T)

Adjoint of trilinear interpolation - distributes residual back to 8 grid points.

# Arguments
- `gradient`: Gradient array to accumulate into (modified in-place)
- `residual`: Observation residual to distribute
- `lon_grid`: Longitude coordinates
- `lat_grid`: Latitude coordinates
- `lev_grid`: Vertical level coordinates
- `lon_obs`: Observation longitude
- `lat_obs`: Observation latitude
- `lev_obs`: Observation vertical level
"""
function trilinear_interpolate_adjoint!(gradient::AbstractArray{T,3},
                                       residual::T,
                                       lon_grid::AbstractVector{T},
                                       lat_grid::AbstractVector{T},
                                       lev_grid::AbstractVector{T},
                                       lon_obs::T, lat_obs::T, lev_obs::T) where T<:AbstractFloat
    weights = compute_trilinear_weights(lon_grid, lat_grid, lev_grid,
                                       lon_obs, lat_obs, lev_obs)

    if !weights.in_bounds
        return nothing
    end

    # Distribute residual according to interpolation weights (adjoint operation)
    for (idx, w) in zip(weights.indices, weights.weights)
        if checkbounds(Bool, gradient, idx)
            gradient[idx] += w * residual
        end
    end

    return nothing
end

# Vertical Interpolation with Different Coordinate Systems

"""
    pressure_to_sigma(pressure::T, surface_pressure::T) -> T

Convert pressure level to sigma coordinate.

σ = p / ps
"""
function pressure_to_sigma(pressure::T, surface_pressure::T) where T<:AbstractFloat
    return pressure / surface_pressure
end

"""
    sigma_to_pressure(sigma::T, surface_pressure::T) -> T

Convert sigma coordinate to pressure level.

p = σ × ps
"""
function sigma_to_pressure(sigma::T, surface_pressure::T) where T<:AbstractFloat
    return sigma * surface_pressure
end

"""
    vertical_interpolate(profile::AbstractVector{T},
                        sigma_levels::AbstractVector{T},
                        sigma_obs::T;
                        method::Symbol=:linear) -> T

Perform vertical interpolation on a profile.

# Arguments
- `profile`: Vertical profile values at model levels
- `sigma_levels`: Model sigma levels
- `sigma_obs`: Observation sigma level
- `method`: Interpolation method (:linear, :logarithmic)

# Returns
- Interpolated value at observation level
"""
function vertical_interpolate(profile::AbstractVector{T},
                             sigma_levels::AbstractVector{T},
                             sigma_obs::T;
                             method::Symbol=:linear) where T<:AbstractFloat
    @assert length(profile) == length(sigma_levels) "Profile and levels must match"

    if method == :linear
        k1, k2, wz = find_grid_cell(sigma_levels, sigma_obs)

        if k1 == k2
            return profile[k1]
        else
            return (T(1) - wz) * profile[k1] + wz * profile[k2]
        end

    elseif method == :logarithmic
        return logarithmic_vertical_interpolate(profile, sigma_levels, sigma_obs)

    else
        error("Unknown vertical interpolation method: $method")
    end
end

"""
    logarithmic_vertical_interpolate(profile::AbstractVector{T},
                                     sigma_levels::AbstractVector{T},
                                     sigma_obs::T) -> T

Logarithmic vertical interpolation (appropriate for temperature, pressure-like fields).

Uses log-linear interpolation: ln(p) vs quantity
"""
function logarithmic_vertical_interpolate(profile::AbstractVector{T},
                                         sigma_levels::AbstractVector{T},
                                         sigma_obs::T) where T<:AbstractFloat
    k1, k2, _ = find_grid_cell(sigma_levels, sigma_obs)

    if k1 == k2 || sigma_levels[k1] <= eps(T) || sigma_levels[k2] <= eps(T)
        return profile[k1]
    end

    # Logarithmic interpolation
    ln_s1 = log(sigma_levels[k1])
    ln_s2 = log(sigma_levels[k2])
    ln_sobs = log(sigma_obs)

    if abs(ln_s2 - ln_s1) < eps(T)
        return profile[k1]
    end

    wz = (ln_sobs - ln_s1) / (ln_s2 - ln_s1)

    return (T(1) - wz) * profile[k1] + wz * profile[k2]
end

"""
    vertical_interpolate_adjoint!(gradient::AbstractVector{T}, residual::T,
                                  sigma_levels::AbstractVector{T},
                                  sigma_obs::T;
                                  method::Symbol=:linear)

Adjoint of vertical interpolation.

# Arguments
- `gradient`: Gradient vector to accumulate into (modified in-place)
- `residual`: Observation residual to distribute
- `sigma_levels`: Model sigma levels
- `sigma_obs`: Observation sigma level
- `method`: Interpolation method (:linear, :logarithmic)
"""
function vertical_interpolate_adjoint!(gradient::AbstractVector{T},
                                      residual::T,
                                      sigma_levels::AbstractVector{T},
                                      sigma_obs::T;
                                      method::Symbol=:linear) where T<:AbstractFloat
    @assert length(gradient) == length(sigma_levels) "Gradient and levels must match"

    k1, k2, wz = find_grid_cell(sigma_levels, sigma_obs)

    if method == :linear
        # Linear interpolation adjoint
        gradient[k1] += (T(1) - wz) * residual
        if k2 != k1 && k2 <= length(gradient)
            gradient[k2] += wz * residual
        end

    elseif method == :logarithmic
        # Logarithmic interpolation adjoint
        if k1 == k2 || sigma_levels[k1] <= eps(T) || sigma_levels[k2] <= eps(T)
            gradient[k1] += residual
        else
            ln_s1 = log(sigma_levels[k1])
            ln_s2 = log(sigma_levels[k2])
            ln_sobs = log(sigma_obs)

            if abs(ln_s2 - ln_s1) > eps(T)
                wz_log = (ln_sobs - ln_s1) / (ln_s2 - ln_s1)
                gradient[k1] += (T(1) - wz_log) * residual
                gradient[k2] += wz_log * residual
            else
                gradient[k1] += residual
            end
        end

    else
        error("Unknown vertical interpolation method: $method")
    end

    return nothing
end

end # module Interpolation
